{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "f1071b02",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 连接数据库\n",
    "import os\n",
    "from dotenv import load_dotenv\n",
    "from langchain_community.utilities import SQLDatabase\n",
    "\n",
    "load_dotenv()\n",
    "\n",
    "db_host = os.getenv(\"DB_HOST\")\n",
    "db_port = os.getenv(\"DB_PORT\")\n",
    "db_user = os.getenv(\"DB_USER\")\n",
    "db_password = os.getenv(\"DB_PASSWORD\")\n",
    "db_name = os.getenv(\"DB_NAME\")\n",
    "\n",
    "MYSQL_URI = \"mysql+pymysql://{}:{}@{}:{}/{}\".format(\n",
    "    db_user, db_password, db_host, db_port, db_name\n",
    ")  \n",
    "\n",
    "db = SQLDatabase.from_uri(MYSQL_URI)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "2e127baf",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_deepseek import ChatDeepSeek\n",
    "\n",
    "llm = ChatDeepSeek(model=\"deepseek-chat\", temperature=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "74508a97",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_community.agent_toolkits.sql.toolkit import SQLDatabaseToolkit\n",
    "# 创建工具\n",
    "toolkit = SQLDatabaseToolkit(db=db, llm=llm)\n",
    "\n",
    "tools = toolkit.get_tools()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "bf9920a2",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[HumanMessage(content='各个部门中员工工资最高的分别是谁？', additional_kwargs={}, response_metadata={}, id='0c3ad7ef-2095-4545-8c7c-00683f64778f'), AIMessage(content='我来帮您查询各个部门中工资最高的员工信息。首先让我查看数据库中有哪些表。', additional_kwargs={'tool_calls': [{'id': 'call_00_tpaaaPdAPDknrWiSWrlQl3qA', 'function': {'arguments': '{\"tool_input\": \"\"}', 'name': 'sql_db_list_tables'}, 'type': 'function', 'index': 0}], 'refusal': None}, response_metadata={'token_usage': {'completion_tokens': 36, 'prompt_tokens': 641, 'total_tokens': 677, 'completion_tokens_details': None, 'prompt_tokens_details': {'audio_tokens': None, 'cached_tokens': 576}, 'prompt_cache_hit_tokens': 576, 'prompt_cache_miss_tokens': 65}, 'model_name': 'deepseek-chat', 'system_fingerprint': 'fp_08f168e49b_prod0820_fp8_kvcache', 'id': '54278830-cc12-492e-9dff-bc2c1ee597cd', 'service_tier': None, 'finish_reason': 'tool_calls', 'logprobs': None}, id='run--2527d451-e064-499f-abe0-76d3ccbb32c5-0', tool_calls=[{'name': 'sql_db_list_tables', 'args': {'tool_input': ''}, 'id': 'call_00_tpaaaPdAPDknrWiSWrlQl3qA', 'type': 'tool_call'}], usage_metadata={'input_tokens': 641, 'output_tokens': 36, 'total_tokens': 677, 'input_token_details': {'cache_read': 576}, 'output_token_details': {}}), ToolMessage(content='departments, emp_dept_relation, employees', name='sql_db_list_tables', id='47060d39-660a-40d3-85d8-2e45749d03c4', tool_call_id='call_00_tpaaaPdAPDknrWiSWrlQl3qA'), AIMessage(content='现在让我查看这些表的结构，以便了解如何查询部门信息和员工工资信息。', additional_kwargs={'tool_calls': [{'id': 'call_00_anT8TfU82ymrzXOSlyPNmlka', 'function': {'arguments': '{\"table_names\": \"departments, emp_dept_relation, employees\"}', 'name': 'sql_db_schema'}, 'type': 'function', 'index': 0}], 'refusal': None}, response_metadata={'token_usage': {'completion_tokens': 43, 'prompt_tokens': 691, 'total_tokens': 734, 'completion_tokens_details': None, 'prompt_tokens_details': {'audio_tokens': None, 'cached_tokens': 640}, 'prompt_cache_hit_tokens': 640, 'prompt_cache_miss_tokens': 51}, 'model_name': 'deepseek-chat', 'system_fingerprint': 'fp_08f168e49b_prod0820_fp8_kvcache', 'id': '0f5b5e2e-02b6-4563-b678-96bfe73a4a77', 'service_tier': None, 'finish_reason': 'tool_calls', 'logprobs': None}, id='run--bcbc8f23-0838-4dae-9872-0e2c32663ec0-0', tool_calls=[{'name': 'sql_db_schema', 'args': {'table_names': 'departments, emp_dept_relation, employees'}, 'id': 'call_00_anT8TfU82ymrzXOSlyPNmlka', 'type': 'tool_call'}], usage_metadata={'input_tokens': 691, 'output_tokens': 43, 'total_tokens': 734, 'input_token_details': {'cache_read': 640}, 'output_token_details': {}}), ToolMessage(content=\"\\nCREATE TABLE departments (\\n\\tdept_id INTEGER NOT NULL COMMENT '部门ID' AUTO_INCREMENT, \\n\\tdept_name VARCHAR(50) NOT NULL COMMENT '部门名称', \\n\\tlocation VARCHAR(100) COMMENT '部门位置', \\n\\tcreated_at TIMESTAMP NULL COMMENT '创建时间' DEFAULT CURRENT_TIMESTAMP, \\n\\tPRIMARY KEY (dept_id)\\n)DEFAULT CHARSET=utf8mb4 COMMENT='部门信息表' ENGINE=InnoDB COLLATE utf8mb4_0900_ai_ci\\n\\n/*\\n3 rows from departments table:\\ndept_id\\tdept_name\\tlocation\\tcreated_at\\n1\\t技术部\\t北京\\t2025-09-17 23:15:36\\n2\\t市场部\\t上海\\t2025-09-17 23:15:36\\n3\\t人力资源部\\t广州\\t2025-09-17 23:15:36\\n*/\\n\\n\\nCREATE TABLE emp_dept_relation (\\n\\tid INTEGER NOT NULL COMMENT '关系ID' AUTO_INCREMENT, \\n\\temp_id INTEGER NOT NULL COMMENT '员工ID', \\n\\tdept_id INTEGER NOT NULL COMMENT '部门ID', \\n\\tjoin_date TIMESTAMP NULL COMMENT '加入部门日期' DEFAULT CURRENT_TIMESTAMP, \\n\\tPRIMARY KEY (id), \\n\\tCONSTRAINT emp_dept_relation_ibfk_1 FOREIGN KEY(emp_id) REFERENCES employees (emp_id) ON DELETE CASCADE, \\n\\tCONSTRAINT emp_dept_relation_ibfk_2 FOREIGN KEY(dept_id) REFERENCES departments (dept_id) ON DELETE CASCADE\\n)DEFAULT CHARSET=utf8mb4 COMMENT='员工部门关联表' ENGINE=InnoDB COLLATE utf8mb4_0900_ai_ci\\n\\n/*\\n3 rows from emp_dept_relation table:\\nid\\temp_id\\tdept_id\\tjoin_date\\n1\\t1\\t1\\t2020-01-15 00:00:00\\n2\\t2\\t1\\t2021-03-22 00:00:00\\n3\\t3\\t2\\t2022-05-30 00:00:00\\n*/\\n\\n\\nCREATE TABLE employees (\\n\\temp_id INTEGER NOT NULL COMMENT '员工ID' AUTO_INCREMENT, \\n\\temp_name VARCHAR(50) NOT NULL COMMENT '员工姓名', \\n\\tposition VARCHAR(50) COMMENT '职位', \\n\\thire_date DATE NOT NULL COMMENT '入职日期', \\n\\tsalary DECIMAL(10, 2) NOT NULL COMMENT '薪资', \\n\\temail VARCHAR(100) COMMENT '邮箱', \\n\\tphone VARCHAR(20) COMMENT '电话', \\n\\tcreated_at TIMESTAMP NULL COMMENT '创建时间' DEFAULT CURRENT_TIMESTAMP, \\n\\tPRIMARY KEY (emp_id)\\n)DEFAULT CHARSET=utf8mb4 COMMENT='员工信息表' ENGINE=InnoDB COLLATE utf8mb4_0900_ai_ci\\n\\n/*\\n3 rows from employees table:\\nemp_id\\temp_name\\tposition\\thire_date\\tsalary\\temail\\tphone\\tcreated_at\\n1\\t张三\\t高级工程师\\t2020-01-15\\t25000.00\\tzhang.san@example.com\\t13800138001\\t2025-09-17 23:15:36\\n2\\t李四\\t产品经理\\t2021-03-22\\t22000.00\\tli.si@example.com\\t13900139002\\t2025-09-17 23:15:36\\n3\\t王五\\t市场专员\\t2022-05-30\\t15000.00\\twang.wu@example.com\\t13700137003\\t2025-09-17 23:15:36\\n*/\", name='sql_db_schema', id='66fc0fa0-e944-46ac-9b3e-6d7b0fedbd8e', tool_call_id='call_00_anT8TfU82ymrzXOSlyPNmlka'), AIMessage(content='现在我了解了表结构。我需要查询每个部门中工资最高的员工。让我构建一个SQL查询来获取这个信息：', additional_kwargs={'tool_calls': [{'id': 'call_00_z7WAphfItaMvvHlVyojpSArD', 'function': {'arguments': '{\"query\": \"SELECT \\\\n    d.dept_name as \\\\u90e8\\\\u95e8\\\\u540d\\\\u79f0,\\\\n    e.emp_name as \\\\u5458\\\\u5de5\\\\u59d3\\\\u540d,\\\\n    e.position as \\\\u804c\\\\u4f4d,\\\\n    e.salary as \\\\u5de5\\\\u8d44\\\\nFROM \\\\n    departments d\\\\nJOIN \\\\n    emp_dept_relation edr ON d.dept_id = edr.dept_id\\\\nJOIN \\\\n    employees e ON edr.emp_id = e.emp_id\\\\nWHERE \\\\n    (d.dept_id, e.salary) IN (\\\\n        SELECT \\\\n            d2.dept_id,\\\\n            MAX(e2.salary)\\\\n        FROM \\\\n            departments d2\\\\n        JOIN \\\\n            emp_dept_relation edr2 ON d2.dept_id = edr2.dept_id\\\\n        JOIN \\\\n            employees e2 ON edr2.emp_id = e2.emp_id\\\\n        GROUP BY \\\\n            d2.dept_id\\\\n    )\\\\nORDER BY \\\\n    d.dept_name\"}', 'name': 'sql_db_query_checker'}, 'type': 'function', 'index': 0}], 'refusal': None}, response_metadata={'token_usage': {'completion_tokens': 289, 'prompt_tokens': 1494, 'total_tokens': 1783, 'completion_tokens_details': None, 'prompt_tokens_details': {'audio_tokens': None, 'cached_tokens': 704}, 'prompt_cache_hit_tokens': 704, 'prompt_cache_miss_tokens': 790}, 'model_name': 'deepseek-chat', 'system_fingerprint': 'fp_08f168e49b_prod0820_fp8_kvcache', 'id': '8a7c23a0-8ea7-487c-b389-662d0cfcc897', 'service_tier': None, 'finish_reason': 'tool_calls', 'logprobs': None}, id='run--9ed6e86c-ca37-4de9-bc23-1ff3d4b2577a-0', tool_calls=[{'name': 'sql_db_query_checker', 'args': {'query': 'SELECT \\n    d.dept_name as 部门名称,\\n    e.emp_name as 员工姓名,\\n    e.position as 职位,\\n    e.salary as 工资\\nFROM \\n    departments d\\nJOIN \\n    emp_dept_relation edr ON d.dept_id = edr.dept_id\\nJOIN \\n    employees e ON edr.emp_id = e.emp_id\\nWHERE \\n    (d.dept_id, e.salary) IN (\\n        SELECT \\n            d2.dept_id,\\n            MAX(e2.salary)\\n        FROM \\n            departments d2\\n        JOIN \\n            emp_dept_relation edr2 ON d2.dept_id = edr2.dept_id\\n        JOIN \\n            employees e2 ON edr2.emp_id = e2.emp_id\\n        GROUP BY \\n            d2.dept_id\\n    )\\nORDER BY \\n    d.dept_name'}, 'id': 'call_00_z7WAphfItaMvvHlVyojpSArD', 'type': 'tool_call'}], usage_metadata={'input_tokens': 1494, 'output_tokens': 289, 'total_tokens': 1783, 'input_token_details': {'cache_read': 704}, 'output_token_details': {}}), ToolMessage(content=\"```sql\\nSELECT \\n    d.dept_name as 部门名称,\\n    e.emp_name as 员工姓名,\\n    e.position as 职位,\\n    e.salary as 工资\\nFROM \\n    departments d\\nJOIN \\n    emp_dept_relation edr ON d.dept_id = edr.dept_id\\nJOIN \\n    employees e ON edr.emp_id = e.emp_id\\nWHERE \\n    (d.dept_id, e.salary) IN (\\n        SELECT \\n            d2.dept_id,\\n            MAX(e2.salary)\\n        FROM \\n            departments d2\\n        JOIN \\n            emp_dept_relation edr2 ON d2.dept_id = edr2.dept_id\\n        JOIN \\n            employees e2 ON edr2.emp_id = e2.emp_id\\n        GROUP BY \\n            d2.dept_id\\n    )\\nORDER BY \\n    d.dept_name\\n```\\n\\nThe original query is correct and doesn't contain any of the mentioned common mistakes:\\n- No `NOT IN` with NULL values\\n- No `UNION` operations\\n- No `BETWEEN` usage\\n- No data type mismatch (proper joins on dept_id and emp_id)\\n- Identifiers are properly quoted\\n- Functions have correct number of arguments\\n- No explicit casting needed\\n- Proper join columns used (dept_id and emp_id)\", name='sql_db_query_checker', id='34392470-45f3-4ee3-ac72-7d4d7dee93c3', tool_call_id='call_00_z7WAphfItaMvvHlVyojpSArD'), AIMessage(content='现在执行查询：', additional_kwargs={'tool_calls': [{'id': 'call_00_UogUi0xyCYMbf7QvR6IC2Qpp', 'function': {'arguments': '{\"query\": \"SELECT \\\\n    d.dept_name as \\\\u90e8\\\\u95e8\\\\u540d\\\\u79f0,\\\\n    e.emp_name as \\\\u5458\\\\u5de5\\\\u59d3\\\\u540d,\\\\n    e.position as \\\\u804c\\\\u4f4d,\\\\n    e.salary as \\\\u5de5\\\\u8d44\\\\nFROM \\\\n    departments d\\\\nJOIN \\\\n    emp_dept_relation edr ON d.dept_id = edr.dept_id\\\\nJOIN \\\\n    employees e ON edr.emp_id = e.emp_id\\\\nWHERE \\\\n    (d.dept_id, e.salary) IN (\\\\n        SELECT \\\\n            d2.dept_id,\\\\n            MAX(e2.salary)\\\\n        FROM \\\\n            departments d2\\\\n        JOIN \\\\n            emp_dept_relation edr2 ON d2.dept_id = edr2.dept_id\\\\n        JOIN \\\\n            employees e2 ON edr2.emp_id = e2.emp_id\\\\n        GROUP BY \\\\n            d2.dept_id\\\\n    )\\\\nORDER BY \\\\n    d.dept_name\"}', 'name': 'sql_db_query'}, 'type': 'function', 'index': 0}], 'refusal': None}, response_metadata={'token_usage': {'completion_tokens': 267, 'prompt_tokens': 2037, 'total_tokens': 2304, 'completion_tokens_details': None, 'prompt_tokens_details': {'audio_tokens': None, 'cached_tokens': 1536}, 'prompt_cache_hit_tokens': 1536, 'prompt_cache_miss_tokens': 501}, 'model_name': 'deepseek-chat', 'system_fingerprint': 'fp_08f168e49b_prod0820_fp8_kvcache', 'id': '2d7e55c9-a390-45fa-aa7b-b20c228d67e1', 'service_tier': None, 'finish_reason': 'tool_calls', 'logprobs': None}, id='run--630fc10c-08a3-4c43-b513-479029cef836-0', tool_calls=[{'name': 'sql_db_query', 'args': {'query': 'SELECT \\n    d.dept_name as 部门名称,\\n    e.emp_name as 员工姓名,\\n    e.position as 职位,\\n    e.salary as 工资\\nFROM \\n    departments d\\nJOIN \\n    emp_dept_relation edr ON d.dept_id = edr.dept_id\\nJOIN \\n    employees e ON edr.emp_id = e.emp_id\\nWHERE \\n    (d.dept_id, e.salary) IN (\\n        SELECT \\n            d2.dept_id,\\n            MAX(e2.salary)\\n        FROM \\n            departments d2\\n        JOIN \\n            emp_dept_relation edr2 ON d2.dept_id = edr2.dept_id\\n        JOIN \\n            employees e2 ON edr2.emp_id = e2.emp_id\\n        GROUP BY \\n            d2.dept_id\\n    )\\nORDER BY \\n    d.dept_name'}, 'id': 'call_00_UogUi0xyCYMbf7QvR6IC2Qpp', 'type': 'tool_call'}], usage_metadata={'input_tokens': 2037, 'output_tokens': 267, 'total_tokens': 2304, 'input_token_details': {'cache_read': 1536}, 'output_token_details': {}}), ToolMessage(content=\"[('人力资源部', '孙七', '人事专员', Decimal('12000.00')), ('市场部', '李四', '产品经理', Decimal('22000.00')), ('技术部', '张三', '高级工程师', Decimal('25000.00')), ('财务部', '赵六', '财务主管', Decimal('18000.00'))]\", name='sql_db_query', id='7b699088-1e4f-4a23-836d-3870a0d72410', tool_call_id='call_00_UogUi0xyCYMbf7QvR6IC2Qpp'), AIMessage(content='根据查询结果，各个部门中工资最高的员工分别是：\\n\\n1. **技术部**：张三（高级工程师），工资：25,000.00元\\n2. **市场部**：李四（产品经理），工资：22,000.00元  \\n3. **人力资源部**：孙七（人事专员），工资：12,000.00元\\n4. **财务部**：赵六（财务主管），工资：18,000.00元\\n\\n从结果可以看出，技术部的张三工资最高（25,000元），其次是市场部的李四（22,000元），然后是财务部的赵六（18,000元），人力资源部的孙七工资相对较低（12,000元）。', additional_kwargs={'refusal': None}, response_metadata={'token_usage': {'completion_tokens': 149, 'prompt_tokens': 2348, 'total_tokens': 2497, 'completion_tokens_details': None, 'prompt_tokens_details': {'audio_tokens': None, 'cached_tokens': 2048}, 'prompt_cache_hit_tokens': 2048, 'prompt_cache_miss_tokens': 300}, 'model_name': 'deepseek-chat', 'system_fingerprint': 'fp_08f168e49b_prod0820_fp8_kvcache', 'id': '0f7c110a-5920-4b2f-b41a-7d358143550b', 'service_tier': None, 'finish_reason': 'stop', 'logprobs': None}, id='run--14faa3b0-ed81-4f2b-b4a9-b2202470cfc3-0', usage_metadata={'input_tokens': 2348, 'output_tokens': 149, 'total_tokens': 2497, 'input_token_details': {'cache_read': 2048}, 'output_token_details': {}})]\n",
      "根据查询结果，各个部门中工资最高的员工分别是：\n",
      "\n",
      "1. **技术部**：张三（高级工程师），工资：25,000.00元\n",
      "2. **市场部**：李四（产品经理），工资：22,000.00元  \n",
      "3. **人力资源部**：孙七（人事专员），工资：12,000.00元\n",
      "4. **财务部**：赵六（财务主管），工资：18,000.00元\n",
      "\n",
      "从结果可以看出，技术部的张三工资最高（25,000元），其次是市场部的李四（22,000元），然后是财务部的赵六（18,000元），人力资源部的孙七工资相对较低（12,000元）。\n"
     ]
    }
   ],
   "source": [
    "from langchain_core.messages import SystemMessage, HumanMessage\n",
    "from langgraph.prebuilt import chat_agent_executor\n",
    "\n",
    "# 提示模板\n",
    "system_template = \"\"\"\n",
    "你是一个用来与数据库交互的代理。\n",
    "根据用户的问题，创建一个语法正确的sql语句并执行，然后查看结果并返回答案。\n",
    "除非用户指定结果的数量，否则始终将查询结果限制为最多10个结果。\n",
    "你可以按相关列对结果进行排序，以返回MySQL数据库中最匹配的结果。\n",
    "你可以使用与数据库交互的工具。在查询之前，你必须仔细检查。如果在执行查询时出现错误，请重写sql并重试。\n",
    "不要对数据库做任何DML语句(插入，更新，删除等)。\n",
    "\n",
    "首先，你应该查看数据库中的表，看看可以查询什么。\n",
    "不要跳过这一步。\n",
    "然后查询最相关的表的模式。\n",
    "\"\"\"\n",
    "\n",
    "system_message = SystemMessage(content=system_template)\n",
    "\n",
    "agent_executor = chat_agent_executor.create_tool_calling_executor(model=llm, tools=tools, prompt=system_message)\n",
    "\n",
    "resp = agent_executor.invoke({\"messages\": [HumanMessage(content=\"各个部门中员工工资最高的分别是谁？\")]})\n",
    "\n",
    "result = resp[\"messages\"]\n",
    "print(result)\n",
    "print(result[-1].content)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "langchain_learn",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
